from langchain_ollama import OllamaLLM
from langchain_community.llms import Ollama
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate
from langchain_core.output_parsers import PydanticOutputParser
from pydantic import BaseModel, Field
from typing import List


# llm = OllamaLLM(
#     model="deepseek-r1:14b",  # 与 ollama list 中的模型名一致
#     base_url="http://localhost:11434",  # Ollama 服务地址
#     temperature=0.3,          # 控制随机性 (0-1)
#     num_predict=2048,         # 最大生成长度
#     top_k=50,                 # 采样参数
#     top_p=0.9,                # 采样参数
#     stop=["<|EOT|>"],         # DeepSeek 的停止标记
#     num_gpu=99,               # 使用全部 GPU 层加速
#     num_thread=8,             # CPU 线程数 (4090 可设 8-12)
# )

from langchain_ollama import ChatOllama
# 1. 初始化 Ollama 模型连接
llm = ChatOllama(model="deepseek-r1:14b", validate_model_on_init=True)

# 1. 定义脚本结构 (Pydantic 模型)
class ShortVideoScript(BaseModel):
    """抖音短视频脚本结构"""
    hook: str = Field(description="黄金开头：5秒内吸引眼球的强冲击力内容", max_length=50)
    setup: str = Field(description="背景铺垫：交代场景和问题，引发共鸣", max_length=100)
    core: List[str] = Field(description="核心内容：分步骤或分要点呈现核心信息（至少3点）")
    emotion: str = Field(description="情绪升华：激发用户情感共鸣（如感动、好奇、愤怒）")
    cta: str = Field(description="行动引导：明确引导用户点赞、关注或评论")
    easter_egg: str = Field(description="结尾彩蛋：意外反转或隐藏福利，提升完播率")

# 2. 配置输出解析器
output_parser = PydanticOutputParser(pydantic_object=ShortVideoScript)

# 3. 构建提示模板 (整合系统指令与用户输入)
system_template = SystemMessagePromptTemplate.from_template(
    "你是一名资深抖音内容策划专家，擅长用{style}风格创作爆款短视频。"
    "严格按以下模块生成脚本：\n"
    "1. 黄金开头（Hook）\n2. 背景铺垫（Setup）\n3. 核心内容（Core）\n"
    "4. 情绪升华（Emotion）\n5. 行动引导（CTA）\n6. 结尾彩蛋（Easter Egg）\n\n"
    "{format_instructions}"
)
prompt = ChatPromptTemplate.from_messages([
    system_template,
    ("human", "主题方向：{topic}")
])

# 4. 组装执行链
chain = prompt | llm | output_parser

# 5. 调用示例
result = chain.invoke({
    "topic": "职场塑料姐妹花",
    "style": "犀利幽默",
    "format_instructions": output_parser.get_format_instructions()
})

print(result)